import os
from volcenginesdkarkruntime import Ark
import datetime
import requests
import hmac
import hashlib

class VolcEngine:
    def __init__(self, api_key):
        self.api_key = api_key
    
    def recognize_image(self, prompt, image_path):
        # 识别图片
        client = Ark(
            api_key=os.environ.get(self.api_key),
        )

        completion = client.bot_chat.completions.create(
            model="Doubao-vision-pro-32k",
            messages = [
                {
                    "role": "user", 
                    "content": [ 
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url", 
                            "image_url": {"url":  image_path}
                        },
                    ],
                }
            ],
        )
        return completion.choices[0].message.content
    
    
    def redraw_style(self, prompt, style, image_path):
        # 风格重绘
        access_key_id = 'example_access_key_id'
        secret_access_key = 'example_access_key'
        region = 'cn-north-1'
        service = 'cv'
        short_date = datetime.datetime.utcnow().strftime('%Y%m%d')
        x_date = datetime.datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')
        
        headers = {
            'Content-Type': 'application/json',
            'Host': 'example.com',
            'X-Date': x_date,
        }
        
        signed_headers = 'content-type;host;x-date'
        canonical_request = f"POST\n/\n\ncontent-type:application/json\nhost:example.com\nx-date:{x_date}\n\n{signed_headers}\n{hashlib.sha256(''.encode('utf-8')).hexdigest()}"
        signature = self.generate_signature(access_key_id, secret_access_key, short_date, region, service, signed_headers, canonical_request)
        
        headers['Authorization'] = f"AWS4-HMAC-SHA256 Credential={access_key_id}/{short_date}/{region}/{service}/aws4_request, SignedHeaders={signed_headers}, Signature={signature}"
        
        body = {
            'Action': 'ImageStyleConversion',
            'Version': '2020-08-26',
            'type': style,
            'image_url': image_path
        }
        
        response = requests.post('https://example.com', headers=headers, json=body)
        return response.json()
    
    def generate_signature(access_key_id, secret_access_key, short_date, region, service, signed_headers, canonical_request):
        string_to_sign = f"AWS4-HMAC-SHA256\n{short_date}T000000Z\n{short_date}/{region}/{service}/aws4_request\n{hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()}"
        date_key = hmac.new(('AWS4' + secret_access_key).encode('utf-8'), short_date.encode('utf-8'), hashlib.sha256).digest()
        date_region_key = hmac.new(date_key, region.encode('utf-8'), hashlib.sha256).digest()
        date_region_service_key = hmac.new(date_region_key, service.encode('utf-8'), hashlib.sha256).digest()
        signing_key = hmac.new(date_region_service_key, b'aws4_request', hashlib.sha256).digest()
        signature = hmac.new(signing_key, string_to_sign.encode('utf-8'), hashlib.sha256).hexdigest()
        return signature
    
    